from lxml import etree


class GEN_Annotations:
    def __init__(self, filename):
        self.root = etree.Element("annotation")

        child1 = etree.SubElement(self.root, "folder")
        child1.text = "VOC2007"

        child2 = etree.SubElement(self.root, "filename")
        child2.text = filename

        child3 = etree.SubElement(self.root, "source")

        child4 = etree.SubElement(child3, "annotation")
        child4.text = "PASCAL VOC2007"
        child5 = etree.SubElement(child3, "database")
        child5.text = "Unknown"

    def set_size(self, witdh, height, channel):
        size = etree.SubElement(self.root, "size")
        widthn = etree.SubElement(size, "width")
        widthn.text = str(witdh)
        heightn = etree.SubElement(size, "height")
        heightn.text = str(height)
        channeln = etree.SubElement(size, "depth")
        channeln.text = str(channel)

    def savefile(self, filename):
        tree = etree.ElementTree(self.root)
        tree.write(filename, pretty_print=True, xml_declaration=False, encoding='utf-8')

    def add_pic_attr(self, label, xmin, ymin, xmax, ymax):
        object = etree.SubElement(self.root, "object")
        namen = etree.SubElement(object, "name")
        namen.text = label
        truncated = etree.SubElement(object, "truncated")
        truncated.text = str(0)
        difficult = etree.SubElement(object, "difficult")
        difficult.text = str(0)
        bndbox = etree.SubElement(object, "bndbox")
        xminn = etree.SubElement(bndbox, "xmin")
        xminn.text = str(xmin)
        yminn = etree.SubElement(bndbox, "ymin")
        yminn.text = str(ymin)
        xmaxn = etree.SubElement(bndbox, "xmax")
        xmaxn.text = str(xmax)
        ymaxn = etree.SubElement(bndbox, "ymax")
        ymaxn.text = str(ymax)



def readTagFromFile(filePath):
    """
    从xml中读取标注数组
    :param filePath:  xml路径
    :return: [{},{}]
    """
    result = []
    xmlfile = etree.parse(filePath)
    root = xmlfile.getroot()
    itemList = root.findall("object", "")
    for item in itemList:
        name = item.find("name", "")
        nameText = name.text
        bndbox = item.find("bndbox", "")
        xmin = bndbox.find("xmin", "")
        xminText = int(xmin.text)
        ymin = bndbox.find("ymin", "")
        yminText = int(ymin.text)
        xmax = bndbox.find("xmax", "")
        xmaxText = int(xmax.text)
        ymax = bndbox.find("ymax", "")
        ymaxText = int(ymax.text)
        result.append({"name": nameText, "xmin": xminText, "ymin": yminText, "xmax": xmaxText, "ymax": ymaxText})
    return result


if __name__ == '__main__':
    pass
